from BitTornado.bencode import bencode
from hashlib import sha1

class MetadataHandler(object):
    """
    Handles Magnet URI downloads
    """
    def __init__(self, rawmetadata, infohash, errorfunc, feedbackfunc = None, finishfunc = None):
        self.rawmetadata = rawmetadata      # Original metainfo info section
        self.infohash = infohash            # Metadata infohash
        self.errorfunc = errorfunc          # Error function    (message)
        self.feedbackfunc = feedbackfunc    # Feedback function (total_size, recieved_size)
        self.finishfunc = finishfunc        # Finish function   (metainfo)

        self.metadata = []                  # Buffer
        self.requested_metadata = []        # How many times each meatdata block has been requested
        self.metadata_progress = 0          # Progress
        self.metadata_size = 0              # Total size
        
        self.hasMetadata = bool(self.rawmetadata)

    def HasMetadata(self):
        """ Returns true in case we have a complete metadata """
        return self.hasMetadata

    def getLength(self):
        """ Get the total size of the metadata """
        return len(self.rawmetadata)
    
    def on_received_metadata(self, data, size, piece, total_size):
        """ New piece of metadata arrived """
        
        print "on_received_metadata", size, piece, total_size
        
        if self.hasMetadata:
            return False

        if total_size <= 0 or total_size > 500*1024:
            return False

        # set metadata size
        if not self.metadata_size:
            if total_size <= 0 or total_size > 500*1024:
                return False
            self.on_metadata_size(total_size)
            
        # verify piece position
        if piece < 0 or piece >= len(self.requested_metadata):
            return False
        
        # set data
        self.metadata[piece] = data

        # feedback
        self.on_metadata_progress(size)

        # mark this piece has 'have'
        self.requested_metadata[piece] = -1

        # check if we're done
        have_all = len(filter(lambda x:x!=-1, self.requested_metadata)) == 0
        if not have_all:
            return False
        
        # Done! Hash it and compare to the original infohash
        self.rawmetadata = ''.join(self.metadata)
        if sha1(self.rawmetadata).digest() != self.infohash:
            self.on_metadata_size(self.metadata_size) # reset :(
            return False
        if self.finishfunc:
            self.finishfunc(self.rawmetadata)
        self.hasMetadata = True
        print "We got the metadata"
        return True
    
    def on_recieved_nodata(self, piece):
        print "on_recieved_nodata", piece
        if self.requested_metadata[piece] <= 0:
            return
        self.requested_metadata[piece] -= 1
        
    def on_recieved_nodata_timeout(self, piece):
        print "on_recieved_nodata_timeout", piece
        if self.requested_metadata[piece] <= 0:
            return
        self.requested_metadata[piece] -= 1
    
    def on_metadata_size(self, size):
        """ Got info about the total size """
        if (type(size) != int or size <= 0 or size > 500 * 1024):
            return
        self.metadata_size = size
        num = ((size + (16*1024) - 1) / (16*1024))
        self.metadata = [""] * num
        self.requested_metadata = [0] * num
        self.on_metadata_progress(0)
        
    def on_metadata_progress(self, received):
        """ give some feedback on progress """
        self.metadata_progress += received;
        if self.feedbackfunc:
            self.feedbackfunc(self.metadata_size, self.metadata_progress)
        print "on_metadata_progress", self.metadata_progress, "/", self.metadata_size

    def write_metadata_packet(self, mType, piece):
        """ write a metadata packet """
        msg = {"msg_type": mType, "piece": piece}
        metadata = ""

        if mType == 1:
            msg["total_size"] = len(self.rawmetadata)
            offset = piece * 16 * 1024
            metadata_piece_size = min(len(self.rawmetadata) - offset, 16 * 1024)
            metadata = self.rawmetadata[offset:offset+metadata_piece_size]
        elif mType == 0:
            self.requested_metadata[piece] += 1

        msg = bencode(msg)
        msg += metadata
        
        return msg

    def next(self):
        """ next piece to be requested """
        # Ask for piece 0
        if not self.requested_metadata:
            self.requested_metadata = [0]
            return 0
            
        for piece in range(0, len(self.requested_metadata)):
            if self.requested_metadata[piece] == 0:
                return piece
        return None
